import numpy as np
import os
import tensorflow as tf
from tensorflow import keras

os.environ['CUDA_VISIBLE_DEVICES'] = '/gpu:0'


def create_vision_net():
    model = keras.models.Sequential()
    model.add(keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(5, 5, 1)))
    model.add(keras.layers.Conv2D(64, (3, 3), activation='relu', input_shape=(3, 3, 1)))
    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(32))
    model.add(keras.layers.Dense(16))
    return model


def create_position_net():
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(8, input_shape=(2,)))
    model.add(keras.layers.Dense(16))
    return model


def create_fusion_net():
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(64, input_shape=(32,)))
    model.add(keras.layers.Dense(128))
    model.add(keras.layers.Dense(5))
    return model


def main():
    vision_net = create_vision_net()
    position_net = create_position_net()
    fusion_net = create_fusion_net()
    vision = tf.zeros((1, 5, 5))
    position = tf.zeros((1, 2))
    vision_out = vision_net(vision)
    position_out = position_net(position)
    fusion_input = tf.concat([vision_out, position_out], axis=1)
    fusion_output = fusion_net(fusion_input)
    print(fusion_output)


if __name__ == '__main__':
    main()
